#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tkinter Desktop GUI Database Viewer
Modern desktop application interface with data browsing and query support
"""

import tkinter as tk
from tkinter import ttk, messagebox, scrolledtext
import pandas as pd
from tkinter import filedialog
import sys
import os

# 数据库连接
try:
    import pg8000 as pg
    DB_DRIVER = 'pg8000'
except ImportError:
    try:
        import psycopg2 as pg
        DB_DRIVER = 'psycopg2'
    except ImportError:
        messagebox.showerror("Error", "Please install database driver: pip install pg8000 or pip install psycopg2-binary")
        sys.exit(1)

class DatabaseViewerGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("🗄️ PostgreSQL Database Viewer")
        self.root.geometry("1200x800")
        self.root.configure(bg='#f0f0f0')
        
        # 数据库连接
        self.connection = None
        self.cursor = None
        self.driver = DB_DRIVER
        self.current_database = None
        
        # 创建界面
        self.create_widgets()
        
    def create_widgets(self):
        """创建界面组件"""
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(2, weight=1)
        
        # 连接设置框架
        conn_frame = ttk.LabelFrame(main_frame, text="🔌 Database Connection", padding="10")
        conn_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        
        # 连接参数
        ttk.Label(conn_frame, text="Host:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.host_var = tk.StringVar(value="localhost")
        ttk.Entry(conn_frame, textvariable=self.host_var, width=15).grid(row=0, column=1, padx=(0, 10))
        
        ttk.Label(conn_frame, text="Port:").grid(row=0, column=2, sticky=tk.W, padx=(0, 5))
        self.port_var = tk.StringVar(value="5432")
        ttk.Entry(conn_frame, textvariable=self.port_var, width=8).grid(row=0, column=3, padx=(0, 10))
        
        ttk.Label(conn_frame, text="User:").grid(row=1, column=0, sticky=tk.W, padx=(0, 5), pady=(5, 0))
        self.user_var = tk.StringVar(value="postgres")
        ttk.Entry(conn_frame, textvariable=self.user_var, width=15).grid(row=1, column=1, padx=(0, 10), pady=(5, 0))
        
        ttk.Label(conn_frame, text="Password:").grid(row=1, column=2, sticky=tk.W, padx=(0, 5), pady=(5, 0))
        self.password_var = tk.StringVar()
        ttk.Entry(conn_frame, textvariable=self.password_var, show="*", width=15).grid(row=1, column=3, padx=(0, 10), pady=(5, 0))
        
        # 连接按钮
        ttk.Button(conn_frame, text="Connect", command=self.connect_database).grid(row=0, column=4, rowspan=2, padx=(10, 0))
        
        # 状态标签
        self.status_var = tk.StringVar(value="Not Connected")
        ttk.Label(conn_frame, textvariable=self.status_var, foreground="red").grid(row=2, column=0, columnspan=5, pady=(5, 0))
        
        # 数据库选择框架
        db_frame = ttk.LabelFrame(main_frame, text="📊 Database Selection", padding="10")
        db_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Label(db_frame, text="Database:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.db_var = tk.StringVar()
        self.db_combo = ttk.Combobox(db_frame, textvariable=self.db_var, width=30, state="readonly")
        self.db_combo.grid(row=0, column=1, padx=(0, 10))
        
        ttk.Button(db_frame, text="Refresh Databases", command=self.refresh_databases).grid(row=0, column=2, padx=(0, 10))
        ttk.Button(db_frame, text="Connect Selected DB", command=self.connect_selected_database).grid(row=0, column=3)
        
        # 主内容区域 - 使用Notebook
        self.notebook = ttk.Notebook(main_frame)
        self.notebook.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        
        # 表列表标签页
        self.tables_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.tables_frame, text="📋 Table List")
        
        # 表列表
        self.tables_tree = ttk.Treeview(self.tables_frame, columns=("type", "rows"), show="tree headings")
        self.tables_tree.heading("#0", text="Table Name")
        self.tables_tree.heading("type", text="Type")
        self.tables_tree.heading("rows", text="Rows")
        self.tables_tree.column("#0", width=300)
        self.tables_tree.column("type", width=100)
        self.tables_tree.column("rows", width=100)
        
        # 滚动条
        tables_scrollbar = ttk.Scrollbar(self.tables_frame, orient=tk.VERTICAL, command=self.tables_tree.yview)
        self.tables_tree.configure(yscrollcommand=tables_scrollbar.set)
        
        self.tables_tree.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        tables_scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        
        # 表列表按钮
        tables_btn_frame = ttk.Frame(self.tables_frame)
        tables_btn_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(5, 0))
        
        ttk.Button(tables_btn_frame, text="Refresh Tables", command=self.refresh_tables).grid(row=0, column=0, padx=(0, 5))
        ttk.Button(tables_btn_frame, text="View Structure", command=self.view_table_structure).grid(row=0, column=1, padx=(0, 5))
        ttk.Button(tables_btn_frame, text="View Data", command=self.view_table_data).grid(row=0, column=2)
        
        # 表数据标签页
        self.data_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.data_frame, text="📄 Table Data")
        
        # 数据表格
        self.data_tree = ttk.Treeview(self.data_frame)
        data_scrollbar_y = ttk.Scrollbar(self.data_frame, orient=tk.VERTICAL, command=self.data_tree.yview)
        data_scrollbar_x = ttk.Scrollbar(self.data_frame, orient=tk.HORIZONTAL, command=self.data_tree.xview)
        self.data_tree.configure(yscrollcommand=data_scrollbar_y.set, xscrollcommand=data_scrollbar_x.set)
        
        self.data_tree.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        data_scrollbar_y.grid(row=0, column=1, sticky=(tk.N, tk.S))
        data_scrollbar_x.grid(row=1, column=0, sticky=(tk.W, tk.E))
        
        # 数据控制
        data_control_frame = ttk.Frame(self.data_frame)
        data_control_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(5, 0))
        
        ttk.Label(data_control_frame, text="Rows to show:").grid(row=0, column=0, padx=(0, 5))
        self.limit_var = tk.StringVar(value="100")
        ttk.Entry(data_control_frame, textvariable=self.limit_var, width=10).grid(row=0, column=1, padx=(0, 10))
        ttk.Button(data_control_frame, text="Refresh Data", command=self.refresh_data).grid(row=0, column=2, padx=(0, 10))
        ttk.Button(data_control_frame, text="Export CSV", command=self.export_csv).grid(row=0, column=3)
        
        # SQL执行器标签页
        self.query_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.query_frame, text="⚡ SQL Executor")
        
        # 查询文本框
        self.query_text = scrolledtext.ScrolledText(self.query_frame, height=10, width=80)
        self.query_text.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        
        # 查询按钮
        query_btn_frame = ttk.Frame(self.query_frame)
        query_btn_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Button(query_btn_frame, text="Execute SQL", command=self.execute_query).grid(row=0, column=0, padx=(0, 5))
        ttk.Button(query_btn_frame, text="Clear", command=self.clear_query).grid(row=0, column=1, padx=(0, 5))
        ttk.Button(query_btn_frame, text="Sample SQL", command=self.load_sample_query).grid(row=0, column=2, padx=(0, 5))
        ttk.Button(query_btn_frame, text="Begin Transaction", command=self.begin_transaction).grid(row=0, column=3, padx=(0, 5))
        ttk.Button(query_btn_frame, text="Commit", command=self.commit_transaction).grid(row=0, column=4, padx=(0, 5))
        ttk.Button(query_btn_frame, text="Rollback", command=self.rollback_transaction).grid(row=0, column=5)
        
        # 执行状态显示
        status_frame = ttk.Frame(self.query_frame)
        status_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 5))
        
        self.execution_status_var = tk.StringVar(value="Ready to execute SQL")
        ttk.Label(status_frame, textvariable=self.execution_status_var, foreground="blue").grid(row=0, column=0, sticky=tk.W)
        
        # 查询结果
        self.query_result_tree = ttk.Treeview(self.query_frame)
        query_result_scrollbar_y = ttk.Scrollbar(self.query_frame, orient=tk.VERTICAL, command=self.query_result_tree.yview)
        query_result_scrollbar_x = ttk.Scrollbar(self.query_frame, orient=tk.HORIZONTAL, command=self.query_result_tree.xview)
        self.query_result_tree.configure(yscrollcommand=query_result_scrollbar_y.set, xscrollcommand=query_result_scrollbar_x.set)
        
        self.query_result_tree.grid(row=3, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        query_result_scrollbar_y.grid(row=3, column=1, sticky=(tk.N, tk.S))
        query_result_scrollbar_x.grid(row=4, column=0, sticky=(tk.W, tk.E))
        
        # 配置网格权重
        self.tables_frame.columnconfigure(0, weight=1)
        self.tables_frame.rowconfigure(0, weight=1)
        self.data_frame.columnconfigure(0, weight=1)
        self.data_frame.rowconfigure(0, weight=1)
        self.query_frame.columnconfigure(0, weight=1)
        self.query_frame.rowconfigure(0, weight=1)
        self.query_frame.rowconfigure(3, weight=1)
        
        # 绑定事件
        self.tables_tree.bind("<Double-1>", self.on_table_double_click)
    
    def _check_cursor_status(self):
        """检查游标状态是否正常"""
        if not self.connection:
            return False, "Database connection is not established"
        
        if not self.cursor:
            return False, "Database cursor is not initialized"
        
        try:
            # 测试游标是否可用
            self.cursor.execute("SELECT 1")
            self.cursor.fetchone()
            return True, "Cursor is ready"
        except Exception as e:
            return False, f"Cursor is not ready: {e}"
    
    def _reset_cursor(self):
        """重置游标状态"""
        try:
            if self.cursor:
                self.cursor.close()
            self.cursor = self.connection.cursor()
            return True
        except Exception as e:
            print(f"Warning: Failed to reset cursor: {e}")
            return False
        
    def connect_database(self):
        """连接数据库"""
        try:
            host = self.host_var.get()
            port = int(self.port_var.get())
            user = self.user_var.get()
            password = self.password_var.get()
            
            if self.driver == 'pg8000':
                self.connection = pg.connect(
                    host=host,
                    port=port,
                    user=user,
                    password=password,
                    database='postgres'
                )
            else:  # psycopg2
                self.connection = pg.connect(
                    host=host,
                    port=port,
                    user=user,
                    password=password,
                    database='postgres'
                )
            
            self.cursor = self.connection.cursor()
            self.status_var.set("Connected to PostgreSQL")
            self.refresh_databases()
            messagebox.showinfo("Success", "Database connected successfully!")
            
        except Exception as e:
            messagebox.showerror("Connection Failed", f"Database connection failed: {e}")
            self.status_var.set("Connection Failed")
    
    def refresh_databases(self):
        """刷新数据库列表"""
        is_ready, message = self._check_cursor_status()
        if not is_ready:
            messagebox.showwarning("Warning", message)
            return
        
        try:
            self.cursor.execute("""
                SELECT datname FROM pg_database 
                WHERE datistemplate = false 
                AND datname NOT IN ('postgres', 'template0', 'template1')
                ORDER BY datname;
            """)
            databases = [row[0] for row in self.cursor.fetchall()]
            self.db_combo['values'] = databases
            if databases:
                self.db_combo.set(databases[0])
        except Exception as e:
            messagebox.showerror("Error", f"Failed to get database list: {e}")
    
    def connect_selected_database(self):
        """连接选中的数据库"""
        if not self.connection:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        selected_db = self.db_var.get()
        if not selected_db:
            messagebox.showwarning("Warning", "Please select a database")
            return
        
        try:
            # 重新连接到选中的数据库
            if self.driver == 'pg8000':
                self.connection = pg.connect(
                    host=self.host_var.get(),
                    port=int(self.port_var.get()),
                    user=self.user_var.get(),
                    password=self.password_var.get(),
                    database=selected_db
                )
            else:
                self.connection = pg.connect(
                    host=self.host_var.get(),
                    port=int(self.port_var.get()),
                    user=self.user_var.get(),
                    password=self.password_var.get(),
                    database=selected_db
                )
            
            self.cursor = self.connection.cursor()
            self.current_database = selected_db
            self.status_var.set(f"Connected to {selected_db}")
            self.refresh_tables()
            messagebox.showinfo("Success", f"Connected to database: {selected_db}")
            
        except Exception as e:
            messagebox.showerror("Connection Failed", f"Failed to connect to database: {e}")
    
    def refresh_tables(self):
        """刷新表列表"""
        if not self.connection or not self.current_database:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        if not self.cursor:
            messagebox.showerror("Error", "Database cursor is not initialized")
            return
        
        try:
            # 清空现有数据
            for item in self.tables_tree.get_children():
                self.tables_tree.delete(item)
            
            # 获取表信息
            self.cursor.execute("""
                SELECT 
                    table_name,
                    table_type,
                    table_schema
                FROM information_schema.tables 
                WHERE table_schema = 'public'
                ORDER BY table_name;
            """)
            tables = self.cursor.fetchall()
            
            for table_name, table_type, table_schema in tables:
                # 获取表的行数 - 使用新的游标避免状态冲突
                try:
                    # 创建临时游标来避免状态冲突
                    temp_cursor = self.connection.cursor()
                    temp_cursor.execute(f"SELECT COUNT(*) FROM {table_name};")
                    row_count = temp_cursor.fetchone()[0]
                    temp_cursor.close()
                except Exception as e:
                    print(f"Warning: Could not get row count for {table_name}: {e}")
                    row_count = 0
                
                self.tables_tree.insert("", "end", text=table_name, values=(table_type, f"{row_count:,}"))
                
        except Exception as e:
            messagebox.showerror("Error", f"Failed to get table list: {e}")
    
    def on_table_double_click(self, event):
        """双击表名时查看表数据"""
        self.view_table_data()
    
    def view_table_structure(self):
        """查看表结构"""
        selection = self.tables_tree.selection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a table")
            return
        
        table_name = self.tables_tree.item(selection[0])['text']
        
        try:
            self.cursor.execute("""
                SELECT 
                    column_name,
                    data_type,
                    is_nullable,
                    column_default,
                    character_maximum_length
                FROM information_schema.columns 
                WHERE table_name = %s 
                AND table_schema = 'public'
                ORDER BY ordinal_position;
            """, (table_name,))
            
            columns = self.cursor.fetchall()
            
            # 创建新窗口显示表结构
            structure_window = tk.Toplevel(self.root)
            structure_window.title(f"Table Structure: {table_name}")
            structure_window.geometry("800x400")
            
            # 创建表格
            structure_tree = ttk.Treeview(structure_window, columns=("type", "nullable", "default", "length"), show="tree headings")
            structure_tree.heading("#0", text="Column Name")
            structure_tree.heading("type", text="Data Type")
            structure_tree.heading("nullable", text="Nullable")
            structure_tree.heading("default", text="Default")
            structure_tree.heading("length", text="Length")
            
            structure_tree.column("#0", width=200)
            structure_tree.column("type", width=150)
            structure_tree.column("nullable", width=80)
            structure_tree.column("default", width=150)
            structure_tree.column("length", width=80)
            
            for col_name, data_type, nullable, default, max_length in columns:
                structure_tree.insert("", "end", text=col_name, values=(
                    data_type,
                    "Yes" if nullable == 'YES' else "No",
                    default or "-",
                    max_length or "-"
                ))
            
            structure_tree.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
            
        except Exception as e:
            messagebox.showerror("Error", f"Failed to get table structure: {e}")
    
    def view_table_data(self):
        """查看表数据"""
        selection = self.tables_tree.selection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a table")
            return
        
        table_name = self.tables_tree.item(selection[0])['text']
        self.refresh_data(table_name)
    
    def refresh_data(self, table_name=None):
        """刷新表数据"""
        if not table_name:
            selection = self.tables_tree.selection()
            if not selection:
                messagebox.showwarning("警告", "请选择表")
                return
            table_name = self.tables_tree.item(selection[0])['text']
        
        try:
            limit = int(self.limit_var.get())
            
            # 获取列名
            self.cursor.execute("""
                SELECT column_name
                FROM information_schema.columns 
                WHERE table_name = %s 
                AND table_schema = 'public'
                ORDER BY ordinal_position;
            """, (table_name,))
            columns = [row[0] for row in self.cursor.fetchall()]
            
            # 获取数据
            self.cursor.execute(f"SELECT * FROM {table_name} LIMIT %s;", (limit,))
            data = self.cursor.fetchall()
            
            # 清空现有数据
            for item in self.data_tree.get_children():
                self.data_tree.delete(item)
            
            # 设置列
            self.data_tree['columns'] = columns
            self.data_tree['show'] = 'headings'
            
            for col in columns:
                self.data_tree.heading(col, text=col)
                self.data_tree.column(col, width=100)
            
            # 插入数据
            for row in data:
                self.data_tree.insert("", "end", values=row)
            
            # 切换到数据标签页
            self.notebook.select(1)
            
        except Exception as e:
            messagebox.showerror("错误", f"获取表数据失败: {e}")
    
    def execute_query(self):
        """执行SQL语句（支持所有SQL语句类型）"""
        if not self.connection:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        if not self.cursor:
            messagebox.showerror("Error", "Database cursor is not initialized")
            return
        
        query = self.query_text.get("1.0", tk.END).strip()
        if not query:
            messagebox.showwarning("Warning", "Please enter a SQL statement")
            return
        
        try:
            # 更新状态
            self.execution_status_var.set("Executing SQL...")
            self.root.update()
            
            # 检查并重置游标状态
            is_ready, message = self._check_cursor_status()
            if not is_ready:
                self.execution_status_var.set(f"Resetting cursor: {message}")
                self.root.update()
                if not self._reset_cursor():
                    raise Exception("Failed to reset cursor")
            
            # 执行SQL语句
            self.cursor.execute(query)
            
            # 判断SQL语句类型
            query_upper = query.upper().strip()
            
            if query_upper.startswith('SELECT') or query_upper.startswith('WITH'):
                # SELECT查询 - 显示结果
                results = self.cursor.fetchall()
                
                if not results:
                    self.execution_status_var.set("Query executed successfully - No results returned")
                    messagebox.showinfo("Result", "Query executed successfully but returned no results")
                    return
                
                # 获取列名
                columns = [desc[0] for desc in self.cursor.description] if hasattr(self.cursor, 'description') else []
                
                # 清空现有数据
                for item in self.query_result_tree.get_children():
                    self.query_result_tree.delete(item)
                
                # 设置列
                self.query_result_tree['columns'] = columns
                self.query_result_tree['show'] = 'headings'
                
                for col in columns:
                    self.query_result_tree.heading(col, text=col)
                    self.query_result_tree.column(col, width=100)
                
                # 插入数据
                for row in results:
                    self.query_result_tree.insert("", "end", values=row)
                
                self.execution_status_var.set(f"Query executed successfully - {len(results)} rows returned")
                
            else:
                # 非SELECT语句 - 显示影响的行数
                affected_rows = self.cursor.rowcount
                
                # 清空结果表格
                for item in self.query_result_tree.get_children():
                    self.query_result_tree.delete(item)
                
                # 设置单列显示执行结果
                self.query_result_tree['columns'] = ("result",)
                self.query_result_tree['show'] = 'headings'
                self.query_result_tree.heading("result", text="Execution Result")
                self.query_result_tree.column("result", width=400)
                
                # 插入执行结果
                result_text = f"SQL statement executed successfully"
                if affected_rows >= 0:
                    result_text += f" - {affected_rows} rows affected"
                
                self.query_result_tree.insert("", "end", values=(result_text,))
                
                self.execution_status_var.set(result_text)
                
                # 刷新表列表（可能表结构已改变）
                self.refresh_tables()
            
            # 切换到SQL执行器标签页
            self.notebook.select(2)
            
        except Exception as e:
            error_msg = f"SQL execution failed: {e}"
            self.execution_status_var.set(error_msg)
            
            # 清空结果表格并显示错误
            for item in self.query_result_tree.get_children():
                self.query_result_tree.delete(item)
            
            self.query_result_tree['columns'] = ("error",)
            self.query_result_tree['show'] = 'headings'
            self.query_result_tree.heading("error", text="Error Message")
            self.query_result_tree.column("error", width=400)
            self.query_result_tree.insert("", "end", values=(error_msg,))
            
            messagebox.showerror("Error", error_msg)
    
    def clear_query(self):
        """清空查询"""
        self.query_text.delete("1.0", tk.END)
    
    def begin_transaction(self):
        """开始事务"""
        if not self.connection:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        try:
            self.cursor.execute("BEGIN;")
            self.execution_status_var.set("Transaction started")
            messagebox.showinfo("Transaction", "Transaction started successfully")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to start transaction: {e}")
    
    def commit_transaction(self):
        """提交事务"""
        if not self.connection:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        try:
            self.cursor.execute("COMMIT;")
            self.execution_status_var.set("Transaction committed")
            messagebox.showinfo("Transaction", "Transaction committed successfully")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to commit transaction: {e}")
    
    def rollback_transaction(self):
        """回滚事务"""
        if not self.connection:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        try:
            self.cursor.execute("ROLLBACK;")
            self.execution_status_var.set("Transaction rolled back")
            messagebox.showinfo("Transaction", "Transaction rolled back successfully")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to rollback transaction: {e}")
    
    def load_sample_query(self):
        """加载示例查询"""
        if not self.current_database:
            messagebox.showwarning("Warning", "Please connect to database first")
            return
        
        # 获取当前数据库的表信息，生成通用示例查询
        try:
            # 获取所有表名
            self.cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
                ORDER BY table_name
                LIMIT 5;
            """)
            tables = [row[0] for row in self.cursor.fetchall()]
            
            if not tables:
                sample_query = """
-- No tables found in public schema
-- Try connecting to a different database or check permissions
SELECT 'No tables available' as message;
                """
            else:
                # 生成基于实际表结构的通用示例查询
                sample_query = self._generate_sample_queries(tables)
                
        except Exception as e:
            # 如果无法获取表信息，使用基本查询
            sample_query = """
-- Basic Database Queries
-- Show all tables
SELECT table_name, table_type
FROM information_schema.tables
WHERE table_schema = 'public'
ORDER BY table_name;

-- Show table structure (replace 'your_table' with actual table name)
SELECT column_name, data_type, is_nullable
FROM information_schema.columns
WHERE table_name = 'your_table'
AND table_schema = 'public'
ORDER BY ordinal_position;

-- Count records in a table (replace 'your_table' with actual table name)
SELECT COUNT(*) as total_records FROM your_table;
            """
        
        self.query_text.delete("1.0", tk.END)
        self.query_text.insert("1.0", sample_query)
    
    def _generate_sample_queries(self, tables):
        """根据实际表结构生成示例查询"""
        queries = ["-- Universal Database Sample Queries", ""]
        
        # 1. 显示所有表
        queries.append("-- 1. Show all tables")
        queries.append("SELECT table_name, table_type")
        queries.append("FROM information_schema.tables")
        queries.append("WHERE table_schema = 'public'")
        queries.append("ORDER BY table_name;")
        queries.append("")
        
        # 2. 为每个表生成智能查询
        for i, table in enumerate(tables[:3]):  # 只处理前3个表
            queries.append(f"-- 2.{i+1} Queries for table '{table}'")
            
            # 表结构查询
            queries.append(f"-- Show structure of {table}")
            queries.append(f"SELECT column_name, data_type, is_nullable")
            queries.append(f"FROM information_schema.columns")
            queries.append(f"WHERE table_name = '{table}'")
            queries.append(f"AND table_schema = 'public'")
            queries.append(f"ORDER BY ordinal_position;")
            queries.append("")
            
            # 记录数查询
            queries.append(f"-- Count records in {table}")
            queries.append(f"SELECT COUNT(*) as total_records FROM {table};")
            queries.append("")
            
            # 前几条记录查询
            queries.append(f"-- Show first 10 records from {table}")
            queries.append(f"SELECT * FROM {table} LIMIT 10;")
            queries.append("")
            
            # 尝试生成智能查询
            smart_queries = self._generate_smart_queries(table)
            if smart_queries:
                queries.extend(smart_queries)
                queries.append("")
        
        # 3. 通用分析查询
        queries.append("-- 3. Database analysis queries")
        queries.append("-- Get table sizes and record counts")
        queries.append("SELECT ")
        queries.append("    schemaname,")
        queries.append("    tablename,")
        queries.append("    attname,")
        queries.append("    n_distinct,")
        queries.append("    correlation")
        queries.append("FROM pg_stats")
        queries.append("WHERE schemaname = 'public'")
        queries.append("ORDER BY tablename, attname;")
        queries.append("")
        
        # 4. 数据质量检查
        queries.append("-- 4. Data quality checks")
        queries.append("-- Find columns with NULL values")
        queries.append("-- (Replace 'table_name' and 'column_name' with actual names)")
        queries.append("-- SELECT column_name, COUNT(*) as null_count")
        queries.append("-- FROM table_name")
        queries.append("-- WHERE column_name IS NULL")
        queries.append("-- GROUP BY column_name;")
        queries.append("")
        
        # 5. 性能分析
        queries.append("-- 5. Performance analysis")
        queries.append("-- Show table sizes")
        queries.append("SELECT ")
        queries.append("    schemaname,")
        queries.append("    tablename,")
        queries.append("    pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) as size")
        queries.append("FROM pg_tables")
        queries.append("WHERE schemaname = 'public'")
        queries.append("ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC;")
        queries.append("")
        
        return "\n".join(queries)
    
    def _generate_smart_queries(self, table_name):
        """为特定表生成智能查询"""
        try:
            # 获取表的列信息
            self.cursor.execute("""
                SELECT column_name, data_type, is_nullable
                FROM information_schema.columns
                WHERE table_name = %s AND table_schema = 'public'
                ORDER BY ordinal_position;
            """, (table_name,))
            columns = self.cursor.fetchall()
            
            if not columns:
                return []
            
            queries = []
            column_names = [col[0] for col in columns]
            data_types = [col[1] for col in columns]
            
            # 查找可能的ID列
            id_columns = [col for col in column_names if 'id' in col.lower()]
            # 查找可能的名称列
            name_columns = [col for col in column_names if any(word in col.lower() for word in ['name', 'title', 'label'])]
            # 查找日期列
            date_columns = [col for col in column_names if any(word in col[1].lower() for word in ['date', 'time', 'timestamp'])]
            # 查找数值列
            numeric_columns = [col for col in column_names if any(word in col[1].lower() for word in ['int', 'numeric', 'decimal', 'float'])]
            
            # 生成基于列类型的智能查询
            if id_columns and name_columns:
                queries.append(f"-- Find records with specific {name_columns[0]} (replace 'value')")
                queries.append(f"SELECT * FROM {table_name}")
                queries.append(f"WHERE {name_columns[0]} LIKE '%value%';")
                queries.append("")
            
            if date_columns:
                queries.append(f"-- Recent records by {date_columns[0]}")
                queries.append(f"SELECT * FROM {table_name}")
                queries.append(f"ORDER BY {date_columns[0]} DESC")
                queries.append(f"LIMIT 10;")
                queries.append("")
            
            if numeric_columns and len(numeric_columns) > 0:
                queries.append(f"-- Statistical analysis of {numeric_columns[0]}")
                queries.append(f"SELECT ")
                queries.append(f"    COUNT(*) as total_count,")
                queries.append(f"    AVG({numeric_columns[0]}) as average_value,")
                queries.append(f"    MIN({numeric_columns[0]}) as min_value,")
                queries.append(f"    MAX({numeric_columns[0]}) as max_value")
                queries.append(f"FROM {table_name};")
                queries.append("")
            
            # 如果有多个列，生成分组查询
            if len(column_names) > 2:
                group_columns = [col for col in column_names if col not in id_columns][:2]
                if group_columns:
                    queries.append(f"-- Group by {group_columns[0]} (if categorical)")
                    queries.append(f"SELECT {group_columns[0]}, COUNT(*) as count")
                    queries.append(f"FROM {table_name}")
                    queries.append(f"GROUP BY {group_columns[0]}")
                    queries.append(f"ORDER BY count DESC;")
                    queries.append("")
            
            return queries
            
        except Exception as e:
            # 如果无法获取列信息，返回空列表
            return []
    
    def export_csv(self):
        """导出CSV"""
        selection = self.tables_tree.selection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a table")
            return
        
        table_name = self.tables_tree.item(selection[0])['text']
        
        try:
            # 获取数据
            self.cursor.execute(f"SELECT * FROM {table_name};")
            data = self.cursor.fetchall()
            
            # 获取列名
            self.cursor.execute("""
                SELECT column_name
                FROM information_schema.columns 
                WHERE table_name = %s 
                AND table_schema = 'public'
                ORDER BY ordinal_position;
            """, (table_name,))
            columns = [row[0] for row in self.cursor.fetchall()]
            
            # 创建DataFrame
            df = pd.DataFrame(data, columns=columns)
            
            # 选择保存位置
            filename = filedialog.asksaveasfilename(
                defaultextension=".csv",
                filetypes=[("CSV files", "*.csv"), ("All files", "*.*")],
                title="Save CSV File"
            )
            
            if filename:
                df.to_csv(filename, index=False, encoding='utf-8')
                messagebox.showinfo("Success", f"Data exported to: {filename}")
            
        except Exception as e:
            messagebox.showerror("Error", f"Export failed: {e}")

def main():
    root = tk.Tk()
    app = DatabaseViewerGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()
